import json
import re
import requests
import time

from db import *
from settings import stocks_url, lei
from pg import cur, conn


def spiderStock(page, url):
    header = {}
    response = requests.get(url % page, header)
    reRes = re.findall(re.compile("[(](.*)[)]"), response.text)
    return json.loads(reRes[0])


def collect_stock(url):
    tmpList = []
    count = 0
    for i in range(1, 143):
        res = spiderStock(i, url)
        stockNameList = res["data"]["diff"]
        for name in stockNameList:
            tmpDict = {}
            if name["f2"] == '-':
                continue
            tmpDict["name"] = name["f14"]
            tmpDict["code"] = name["f12"]
            tmpDict["price"] = name["f20"]
            if name["f13"]:
                tmpDict["prefix"] = "SH"
            else:
                tmpDict["prefix"] = "SZ"
            if str(tmpDict["code"]).startswith("8") or str(tmpDict["code"]).startswith("4"):
                tmpDict["prefix"] = "BJ"

            tmpDict["exist"] = 1
            tmpDict["symbol"] = tmpDict['prefix'].upper() + tmpDict["code"]
            if '退' in tmpDict["name"] or tmpDict["name"] in lei:
                tmpDict["exist"] = 0
            count += 1
            print(tmpDict["symbol"], name["f13"], count)

            # stock_db.update_one(
            #     {"symbol": tmpDict["symbol"]},
            #     {"$set": tmpDict},
            #     upsert=True
            # )
            insert_query = """
            INSERT INTO stock_info (symbol, code, exist, name, prefix, price)
            VALUES (%s, %s, %s, %s, %s, %s)
            ON CONFLICT (symbol) 
            DO UPDATE SET 
                code = EXCLUDED.code,
                exist = EXCLUDED.exist,
                name = EXCLUDED.name,
                prefix = EXCLUDED.prefix,
                price = EXCLUDED.price;
            """
            data = (tmpDict["symbol"], tmpDict["code"] , tmpDict["exist"], tmpDict["name"], tmpDict["prefix"], tmpDict["price"])

            cur.execute(insert_query, data)


        conn.commit()
            
        time.sleep(1)
    return tmpList


if __name__ == "__main__":
    collect_stock(stocks_url)
